#define L1_BUFFER_LEN ::std::size_t(128)

__global__ void kernel_add(float *lhs, float *rhs, float *out, size_t elems_per_thread) {
    tops_dte_ctx_t ctx;
    tops::dte_scope _(ctx); // Initialize DTE context and manage its lifetime

    // Calculate the unique thread index within the grid
    uint32_t const thread_idx {blockIdx.x * blockDim.x + threadIdx.x};
    auto pos_start = thread_idx * elems_per_thread;
    auto pos_end = pos_start + elems_per_thread;

    // Define L1 buffers with proper alignment
    __valigned__ float l1_buf_lhs[L1_BUFFER_LEN];
    __valigned__ float l1_buf_rhs[L1_BUFFER_LEN];
    __valigned__ float l1_buf_out[L1_BUFFER_LEN];

    // Create mdspan for L1 buffers
    tops::mdspan l1_lhs(tops::Private, l1_buf_lhs, L1_BUFFER_LEN);
    tops::mdspan l1_rhs(tops::Private, l1_buf_rhs, L1_BUFFER_LEN);
    tops::mdspan l1_out(tops::Private, l1_buf_out, L1_BUFFER_LEN);

    for (::std::size_t i = pos_start; i < pos_end; i += L1_BUFFER_LEN) {
        // Create mdspan for current chunk in L3
        tops::mdspan l3_lhs(tops::Global, lhs + i, L1_BUFFER_LEN);
        tops::mdspan l3_rhs(tops::Global, rhs + i, L1_BUFFER_LEN);

        // Asynchronously load data from L3 to L1
        tops::memcpy(ctx, l1_lhs, l3_lhs);
        tops::memcpy(ctx, l1_rhs, l3_rhs);

        // Perform vectorized addition on L1 buffers
        for (size_t j = 0; j < L1_BUFFER_LEN; j += tops::vlength<vfloat>()) {
            const auto &v_lhs = tops::vload<vfloat>(l1_buf_lhs + j);
            const auto &v_rhs = tops::vload<vfloat>(l1_buf_rhs + j);
            const auto &v_res = tops::vadd<vfloat>(v_lhs, v_rhs);
            tops::vstore(v_res, l1_buf_out + j);
        }

        // Create mdspan for the output chunk in L3
        tops::mdspan l3_out(tops::Global, out + i, L1_BUFFER_LEN);

        // Asynchronously write result back to L3
        tops::memcpy(ctx, l3_out, l1_out);
    }
}

void GCU_ADD(float * __restrict dev_lhs, float * __restrict dev_rhs, float * __restrict dev_out, const int nr_elems) {
    static const size_t blocks = 2;
    static const size_t threads = 12;
    kernel_add<<<blocks, threads>>>(dev_lhs, dev_rhs, dev_out, nr_elems / (blocks * threads));
}
